package cz.cuni.amis.episodic.dybanem;

import cz.cuni.amis.episodic.bayes.memories.KLBasedMemoryCreator;
import cz.cuni.amis.episodic.bayes.memories.MemoriesUtils;
import cz.cuni.amis.episodic.bayes.memories.MemoryElement;
import cz.cuni.amis.episodic.bayes.memories.MemoryCreator;
import cz.cuni.amis.episodic.bayes.utils.NodeSnapshot;
import cz.cuni.amis.episodic.bayes.utils.chart.AreaChart;
import cz.cuni.amis.episodic.bayes.utils.chart.GraphPaintingDevice;

import java.awt.Color;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.List;
import org.apache.commons.lang.StringUtils;

/**
 *
 * @author ik
 */
public class RetrospectiveMemoriesCreator {

    /**
     * Creates memories based on surprise associated with the mem.
     *
     * @param dbn
     * @param evidence
     * @param targets
     * @return
     * @throws IOException
     */
    public static List<MemoryElement> createMemories(DBN dbn, List<String> evidence, String[] targets, int iterations, MemoryCreator memoryCreator, GraphPaintingDevice device) throws IOException {
        final String evidenceNode = "O";
        // get distribution based on last complete observation
        // this will be the "true" distribution 
        dbn.getNet().clearAllEvidence();
        dbn.setDBNTargetNodes(targets);
        dbn.setDBNEvidence(evidenceNode, evidence);
        dbn.getNet().updateBeliefs();
        NodeSnapshot[] trueSnapshots = dbn.getNodesSnapshot(targets);

        // iteratively add memories to make KL divergence between remembered state and
        // the true distribution minimal
        List<MemoryElement> memories = new ArrayList<>();
        
        for (int i = 0; i < iterations; i++) {
            // clear evidence and set new
            dbn.resetEvidence(memories);
            dbn.getNet().updateBeliefs();
            if(device != null) {
            	NodeSnapshot[] finalSnapshots = dbn.getNodesSnapshot(targets);
            	// TODO refactore
            	// entropy
                drawEntropyOfRecall(finalSnapshots, i, device);
            }

            MemoryElement mem = memoryCreator.createMemory(targets, trueSnapshots, dbn, memories, i);

            memories.add(mem);

            // add it as an evidence to the network and continue
            // TODO this is just optional
            mem.setEvidence(dbn);
            // print memories
            MemoriesUtils.printMemories(memories, System.out);

            
            // print target distribution
            dbn.getNet().updateBeliefs();
            PrintStream out = System.out;
            out.println(StringUtils.join(evidence, ' '));
            out.println(StringUtils.join(dbn.getDBNNodeMostProbableStringValues(evidenceNode), ' '));
            out.println();

        }

        if(device != null) {
        	NodeSnapshot[] finalSnapshots = dbn.getNodesSnapshot(targets);
        	// TODO refactore
        	// entropy
            drawEntropyOfRecall(finalSnapshots, iterations, device);
            // prior probability
        	KLBasedMemoryCreator.drawPriorProbabilityGraphs(finalSnapshots, dbn, iterations, device, memoryCreator.getStrategyName());
        }
        dbn.getNet().clearAllTargets();
        //double[][] transpose = MemChart.transpose(kl);
        return memories;
    }
    
    public static void drawEntropyOfRecall(NodeSnapshot[] nodeSnapshots, int iteration, GraphPaintingDevice device) {
    	for(NodeSnapshot snapshot : nodeSnapshots) {
    		double[] entropy = ProbabilityUtils.entropy(snapshot.getProbability());
    		try {
				device.outputGraph(
				        AreaChart.chart(entropy, Color.green, null, "Entropy"),
				        snapshot.getNode() + "_entropy", iteration);
			} catch (Exception e) {
				e.printStackTrace();
			}
     
    	}
    }
}
